import psycopg2
import json
from tqdm import tqdm

class VBAFile2Database:
    def __init__(self, model_name):
        """
        初始化数据库连接
        :param model_name: 不再使用，保留参数以保持接口兼容
        """
        self.DB_CONFIG = {
            "host": "hadoop102",
            "port": 15432,
            "database": "hirain_vba_usage",
            "user": "postgres",
            "password": "postgres"
        }

    def store_in_db(self, documents, source):
        """存储到PostgreSQL数据库"""
        conn = psycopg2.connect(**self.DB_CONFIG)

        try:
            with conn.cursor() as cur:
                # 创建存储表（不再需要向量字段）
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS vba_docs (
                        id SERIAL PRIMARY KEY,
                        content TEXT,
                        metadata JSONB,
                        collection_name VARCHAR(255)
                    )
                """)

                with tqdm(documents) as pbar:
                    for doc in documents:
                        metadata = {
                            "source": doc['source'],
                            "collection": doc['collection']
                        }

                        metadata = json.dumps(metadata)
                        cur.execute("""
                        INSERT INTO vba_docs (content, metadata, collection_name)
                        VALUES (%s, %s, %s)
                        """, (doc['answer'], metadata, doc['source']))
                        pbar.update()
                conn.commit()
                print(f"成功插入 {len(documents)} 条记录")
        finally:
            conn.close()
